from fastapi import FastAPI
from contextlib import asynccontextmanager
from extends.db import engine, Base, async_session
from models.user import Power
from controllers import create_router
from sqlalchemy import select
from repository import RepositoryExt, RepositoryNotFoundExt
from service import ServiceExtends
from fastapi.responses import JSONResponse


@asynccontextmanager
async def lifespan(app: FastAPI):
    # await startup()
    # await initialize_database()
    yield
    await shutdown()

app = FastAPI(lifespan=lifespan)

@app.exception_handler(RepositoryNotFoundExt)
async def repository_not_found_exception_handler(request, exc: RepositoryNotFoundExt):
    return JSONResponse(status_code=400, content={"code": 0, "msg": "not found: " + exc.msg})

@app.exception_handler(RepositoryExt)
async def repository_not_found_exception_handler(request, exc: RepositoryExt):
    return JSONResponse(status_code=400, content={"code": 0, "msg": "repository error: " + exc.msg})

@app.exception_handler(ServiceExtends)
async def repository_not_found_exception_handler(request, exc: ServiceExtends):
    return JSONResponse(status_code=400, content={"code": 0, "msg": "service error: " + exc.msg})

create_router(app)

async def startup():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)


async def initialize_database():
    async with async_session() as session:
        # 初始化权限
        for table_name in Base.metadata.tables:
            init_power = [
                Power(name=f"get_{table_name}"),
                Power(name=f"create_{table_name}"),
                Power(name=f"update_{table_name}"),
                Power(name=f"delete_{table_name}"),
            ]
            for p in init_power:
                res = await session.execute(select(Power).where(Power.name == p.name))
                if not res.scalar():
                    session.add(p)
            await session.commit()


async def shutdown():
    await engine.dispose()

async def get_db():
    db = async_session()
    try:
        yield db
    finally:
        await db.close()
    

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)